"""
Sensor Array Flow Scheduling - Heuristic Simulation
----------------------------------------------------
- 20x30 sensor grid, 3 moving cars (each covers fixed two rows)
- Cars move along x-axis, covering a 2x3 block at any time
- Flows generated at sensors at specific times
- Heuristic per-second simulation: flows move one hop per second toward a covered node,
  then upload to car if possible, limited by bandwidth constraints.
- Metrics: total/lost data, loss rate, number of completed flows, avg completion delay
- Also outputs per-car receive curves, completion delay histogram, and sample flow traces.

Author: ChatGPT
"""

import math
import random
import networkx as nx
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from collections import deque, defaultdict

# -----------------------------
# Problem parameters
# -----------------------------
M, N = 20, 30          # grid size
T = 90                 # simulation seconds
B_sensor = 10.0        # Mbps per directed link
B_receive = 100.0      # Mbps per car
flow_size = 10.0       # Mb per flow
flow_rate_max = 5.0    # Mb/s per flow instantaneous
produce_times = [0, 30, 60]  # for N even
period_b = 10.0        # period for b(t)

# Rows each car covers (y coordinates, 1-based)
car_rows = [(5,6), (15,16), (25,26)]

# -----------------------------
# Build grid graph
# -----------------------------
G = nx.DiGraph()
for x in range(1, M+1):
    for y in range(1, N+1):
        G.add_node((x,y))

# Add 4-neighbor directed edges
for x in range(1, M+1):
    for y in range(1, N+1):
        for dx, dy in [(1,0), (-1,0), (0,1), (0,-1)]:
            nx_, ny_ = x+dx, y+dy
            if 1 <= nx_ <= M and 1 <= ny_ <= N:
                G.add_edge((x,y), (nx_,ny_), capacity=B_sensor)

# -----------------------------
# Sensor-to-car bandwidth parameters
# -----------------------------
random.seed(0)
B_peak = {node: random.uniform(8.0, 15.0) for node in G.nodes()}  # peak between 8 and 15 Mbps
phi = {node: random.uniform(0, period_b) for node in G.nodes()}   # phase shift

def b_u(node, t):
    """Triangular periodic bandwidth function between 0.5*peak and 1.0*peak."""
    peak = B_peak[node]
    val = ((t + phi[node]) % period_b)
    rel = 1.0 - abs(val - period_b/2) / (period_b/2)  # 0..1
    return 0.5*peak + 0.5*peak*rel

# -----------------------------
# Car movement & coverage
# -----------------------------
def car_x_position(t, car_index):
    frac = (t % 100) / 100.0
    pos = 1 + frac*(M-3)
    return pos

def covered_nodes_of_car(k, t):
    rows = car_rows[k]
    left_col = int(math.floor(car_x_position(t, k)))
    left_col = max(1, min(M-2, left_col))
    cols = [left_col, left_col+1, left_col+2]
    covered = []
    for x in cols:
        for y in rows:
            covered.append((x,y))
    return set(covered)

# -----------------------------
# Flow generation
# -----------------------------
flows = []
flow_id = 0
for node in G.nodes():
    for pt in produce_times:
        if pt < T:
            flows.append({
                'id': flow_id,
                'src': node,
                't_start': pt,
                'size': flow_size,
                'remaining': flow_size,
                'completed': False,
                'completion_time': None,
                'path': [node],  # 添加路径记录，初始化为起始节点
                'path_times': [pt]  # 记录到达每个节点的时间
            })
            flow_id += 1
F = len(flows)
flows_by_id = {f['id']: f for f in flows}

# -----------------------------
# Utility: shortest path BFS
# -----------------------------
def shortest_path_to_any_target(src, targets):
    if src in targets:
        return [src]
    visited = {src}
    dq = deque([(src, [src])])
    while dq:
        u, path = dq.popleft()
        for v in G.successors(u):
            if v in visited:
                continue
            visited.add(v)
            newp = path + [v]
            if v in targets:
                return newp
            dq.append((v, newp))
    return None

# -----------------------------
# Simulation
# -----------------------------
def reset_link_caps():
    return {(u,v): B_sensor for u,v in G.edges()}

total_received = 0.0
total_dropped = 0.0
per_second_car_recv = [[0.0]*T for _ in range(len(car_rows))]
sample_paths = {}

def get_active_flows(t):
    return [f for f in flows if f['t_start'] <= t and f['remaining'] > 1e-9]

for t in range(T):
    link_caps = reset_link_caps()
    car_recv_remaining = {k: B_receive for k in range(len(car_rows))}
    node_to_car_used = defaultdict(float)
    car_covers = {k: covered_nodes_of_car(k, t) for k in range(len(car_rows))}
    active = get_active_flows(t)
    active.sort(key=lambda ff: (ff['t_start'], ff['remaining']))

    for f in active:
        uid, src, rem = f['id'], f['src'], f['remaining']
        if rem <= 1e-9:
            continue

        # Upload directly if in coverage
        covered_pairs = []
        for k, cov in car_covers.items():
            if src in cov:
                bnode = b_u(src, t)
                avail_node = max(0.0, bnode - node_to_car_used[(src,k)])
                avail_car = car_recv_remaining[k]
                covered_pairs.append((k, avail_node, avail_car))

        if covered_pairs:
            covered_pairs.sort(key=lambda x: (-x[2], -x[1]))
            chosen = covered_pairs[0][0]
            allowed = min(rem, flow_rate_max, covered_pairs[0][1], covered_pairs[0][2])
            if allowed > 1e-9:
                f['remaining'] -= allowed
                car_recv_remaining[chosen] -= allowed
                node_to_car_used[(src,chosen)] += allowed
                total_received += allowed
                per_second_car_recv[chosen][t] += allowed
                if f['remaining'] <= 1e-9 and f['completion_time'] is None:
                    f['completed'] = True
                    f['completion_time'] = t
                    f['final_car'] = chosen  # 记录上传到哪辆车
            continue

        # Otherwise, move toward coverage
        possible_targets = set()
        for k, cov in car_covers.items():
            if car_recv_remaining[k] > 1e-9:
                possible_targets |= cov
        if not possible_targets:
            continue
        path = shortest_path_to_any_target(src, possible_targets)
        if not path or len(path) < 2:
            continue
        next_hop = path[1]
        cap = link_caps.get((src,next_hop), 0.0)
        allowed = min(rem, flow_rate_max, cap)
        if allowed <= 1e-9:
            continue
        link_caps[(src,next_hop)] -= allowed
        f['remaining'] -= allowed
        f['src'] = next_hop  # move flow
        # 记录路径和时间
        f['path'].append(next_hop)
        f['path_times'].append(t + 1)  # t+1 表示下一秒到达
        # (We don't model mid-second arrivals; next second this new src will be used)

# -----------------------------
# Paths
# -----------------------------
def export_flow_paths():
    # 准备数据
    all_data = []
    for f in flows:
        flow_data = {
            'flow_id': f['id'],
            'start_time': f['t_start'],
            'completion_time': f['completion_time'],
            'completed': f['completed'],
            'original_size': f['size'],
            'remaining_size': f['remaining'],
            'path_length': len(f['path']),
        }

        # 添加路径坐标
        for i, node in enumerate(f['path']):
            flow_data[f'node_{i}_x'] = node[0]
            flow_data[f'node_{i}_y'] = node[1]
            flow_data[f'node_{i}_time'] = f['path_times'][i]

        all_data.append(flow_data)
    # 转换为DataFrame并导出
    df = pd.DataFrame(all_data)
    # 导出为CSV
    df.to_csv('flow_paths.csv', index=False)
    # 导出为XLSX
    # df.to_excel('flow_paths.xlsx', index=False)
    print(f"已导出 {len(flows)} 条流量的路径数据到 flow_paths.csv")


def export_detailed_paths():
    # 创建一个包含所有流量路径的列表
    path_records = []

    for f in flows:
        for i in range(len(f['path']) - 1):
            # 对每个路径段
            path_records.append({
                'flow_id': f['id'],
                'step': i,
                'from_x': f['path'][i][0],
                'from_y': f['path'][i][1],
                'to_x': f['path'][i + 1][0],
                'to_y': f['path'][i + 1][1],
                'time': f['path_times'][i],
                'next_time': f['path_times'][i + 1]
            })

    # 如果有完成的流量，添加最终上传到车辆的记录
    for f in flows:
        if f['completed'] and hasattr(f, 'final_car'):
            path_records.append({
                'flow_id': f['id'],
                'step': len(f['path']) - 1,
                'from_x': f['path'][-1][0],
                'from_y': f['path'][-1][1],
                'to_x': 'car',
                'to_y': f['final_car'],
                'time': f['path_times'][-1],
                'next_time': f['completion_time']
            })

    # 导出详细路径
    df_detailed = pd.DataFrame(path_records)
    df_detailed.to_csv('detailed_flow_paths.csv', index=False)
    #df_detailed.to_excel('detailed_flow_paths.xlsx', index=False)

    print(f"已导出详细路径数据到 detailed_flow_paths.csv")


export_detailed_paths()
export_flow_paths()
# -----------------------------
# Metrics
# -----------------------------
Q_total = sum(f['size'] for f in flows)
Q_recv = total_received
loss_rate = 1 - (Q_recv / Q_total)
completed = [f['completion_time'] - f['t_start']
             for f in flows if f['completion_time'] is not None]
avg_delay = sum(completed)/len(completed) if completed else None
num_completed = len(completed)

print("Simulation metrics:")
print(f"  Total flows: {F}")
print(f"  Q_total(Mb): {Q_total}")
print(f"  Q_received(Mb): {Q_recv}")
print(f"  Loss rate: {loss_rate:.4f}")
print(f"  Completed flows: {num_completed}")
print(f"  Avg delay (completed): {avg_delay}")

# -----------------------------
# Plots
# -----------------------------
plt.figure(figsize=(10,4))
for k in range(len(car_rows)):
    plt.plot(range(T), per_second_car_recv[k])
plt.xlabel("Time (s)")
plt.ylabel("Received Mb/s")
plt.title("Per-car received rate over time")
plt.legend([f"Car{k}" for k in range(len(car_rows))])
plt.grid(True)
plt.show()

if completed:
    plt.figure(figsize=(6,4))
    plt.hist(completed, bins=20)
    plt.xlabel("Completion delay (s)")
    plt.ylabel("Number of flows")
    plt.title("Completion delay histogram")
    plt.show()
